package com.selftcp;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.SortedMap;
import java.util.Timer;
import java.util.TimerTask;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

// 模拟TCP连接 - 不处理粘包拆包，保留原始特性
public class TCPConnection {
	private final int localPort;
    private final int remotePort;
    private final InetAddress remoteAddr;
    private final DatagramSocket socket;
    
    // 用于接收服务器分发的数据包
    private final LinkedBlockingQueue<TCPacket> receiveQueue = new LinkedBlockingQueue<>();
    
    // 序列号相关
    private long sendSeq;       // 下一个要发送的序列号
    private long recvSeq;       // 期望接收的下一个序列号
    private long lastAck;       // 最后确认的序列号
    
    // 拥塞控制相关
    private int cwnd = 1;       // 拥塞窗口（慢启动开始）
    private int ssthresh = 64;  // 慢启动阈值
    private int rwnd = 64;      // 接收窗口
    private int mss = 50;       // 最大段大小
    
    // 重传机制
    private final Map<Long, TCPacket> unackedPackets = new ConcurrentHashMap<>();
    private final Timer retransmitTimer = new Timer();
    
    // 接收缓冲区 - 用于重组乱序数据包
    private final SortedMap<Long, byte[]> receiveBuffer = new TreeMap<>();
    
    // 发送和接收队列
    private final LinkedBlockingQueue<byte[]> sendQueue = new LinkedBlockingQueue<>();
    private final LinkedBlockingQueue<byte[]> appReceiveQueue = new LinkedBlockingQueue<>();
    
    // 连接状态
    private volatile boolean isConnected = false;
    private volatile boolean isClosing = false;
    
    public TCPConnection(int localPort, InetAddress remoteAddr, int remotePort, DatagramSocket socket) {
        this.localPort = localPort;
        this.remoteAddr = remoteAddr;
        this.remotePort = remotePort;
        this.socket = socket;
        
        // 初始化随机序列号
        this.sendSeq = new Random().nextLong() & 0x7FFFFFFFFFFFFFFFL;
        this.recvSeq = 0;
        this.lastAck = 0;
    }
    
    // 建立连接（三次握手）
    public void connect() throws IOException, InterruptedException {
    	// 发送SYN包，SYN包的全称：Synchronize Packet（同步数据包），作用：同步序列号，为可靠传输 “打底”
    	// 1、当一方（如客户端）发起连接时，会生成一个随机的初始序列号（ISN，Initial Sequence Number），并通过 SYN 包将这个 ISN 发送给对方（如服务器）；
    	// 2、对方收到 SYN 包后，会确认 “已收到你的 ISN”，同时也生成自己的 ISN，通过 “SYN+ACK 包”（同时携带同步和确认标志）将自己的 ISN 发送回去；
    	// 3、发起方再通过 ACK 包确认 “已收到你的 ISN”—— 至此，双方的序列号同步完成，后续传输的数据都会基于这个初始序列号递增，确保顺序和完整性。
        sendSeq = new Random().nextLong() & 0x7FFFFFFFFFFFFFFFL;
        TCPacket synPacket = new TCPacket(localPort, remotePort, sendSeq, 0, true, false, false, rwnd, null);
        System.out.println("三次握手，第一次握手"+synPacket);
        sendPacket(synPacket);
        sendSeq++;
        
        // 启动客户端监听线程
        startClientListener();
        
        // 等待SYN-ACK（从receiveQueue接收）
        while (!isConnected) {
        	TCPacket tcpPacket = receiveQueue.poll(1000, TimeUnit.MILLISECONDS);
            if (tcpPacket == null) continue;
            
            System.out.println("三次握手，第二次握手：" + tcpPacket);
            if (tcpPacket.syn && tcpPacket.ackFlag && tcpPacket.ack == sendSeq) {
            	// 清理未确认的SYN包
                synchronized (unackedPackets) {
                	TCPacket t = unackedPackets.remove(sendSeq - 1); // 初始SYN包的序列号是sendSeq-1
                    System.out.println("删除第一次握手重试：" + (sendSeq - 1) + (t!=null?"成功":"失败"));
                }
                
                recvSeq = tcpPacket.seq + 1;
                lastAck = recvSeq;
                
                // 发送ACK，完成三次握手
                TCPacket ackPacket = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                                false, true, false, rwnd, null);
                sendPacket(ackPacket);
                System.out.println("三次握手，第三次握手：" + ackPacket);
                
                isConnected = true;
                startSender();
                startReceiverProcessor();
                break;
            }
        }
    }
    
    // 客户端专用的Socket监听线程，负责接收所有数据并分发到receiveQueue
    private void startClientListener() {
        new Thread(() -> {
            try {
                while (!isClosing) {
                    byte[] buf = new byte[4096];
                    DatagramPacket packet = new DatagramPacket(buf, buf.length);
                    socket.receive(packet);
                    
                    // 验证数据包来源（确保是目标服务器）
                    if (!packet.getAddress().equals(remoteAddr) || packet.getPort() != remotePort) {
                        continue;
                    }
                    
                    // 解析并分发数据包到receiveQueue
                    TCPacket tcpPacket = TCPacket.fromBytes(packet.getData());
                    dispatchPacket(tcpPacket);
                }
            } catch (IOException e) {
                if (!isClosing) {
                    e.printStackTrace();
                }
            }
        }).start();
    }
    
    // 处理接收到的连接请求（三次握手的服务器端）
    public void handleConnectionRequest(TCPacket synPacket) throws IOException {
    	System.out.println("三次握手，第一次握手：" + synPacket);
        recvSeq = synPacket.seq + 1;
        sendSeq = new Random().nextLong() & 0x7FFFFFFFFFFFFFFFL;
        
        // 发送SYN-ACK
        TCPacket synAckPacket = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                           true, true, false, rwnd, null);
        sendPacket(synAckPacket);
        sendSeq++;
    }
    
    // 完成连接建立（三次握手的最后一步）
    public boolean completeConnection(TCPacket ackPacket) {
        if (ackPacket.ackFlag && ackPacket.ack == sendSeq) {
        	System.out.println("三次握手，第三次握手" + ackPacket);
            isConnected = true;
            startSender();
            startReceiverProcessor();
            return true;
        }
        return false;
    }
    
    // 发送数据
    public void send(byte[] data) throws InterruptedException {
        if (!isConnected) {
            throw new IllegalStateException("Not connected");
        }
        
        // 按MSS拆分数据
        int offset = 0;
        while (offset < data.length) {
            int length = Math.min(data.length - offset, mss);
            byte[] chunk = new byte[length];
            System.arraycopy(data, offset, chunk, 0, length);
            sendQueue.put(chunk);
            offset += length;
        }
    }
    
    // 接收数据
    public byte[] receive() throws InterruptedException {
        return appReceiveQueue.take();
    }
    
    // 关闭连接
    public void close() throws IOException {
        if (isClosing || !isConnected) return;
        
        isClosing = true;
        // 发送FIN包
        TCPacket finPacket = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                        false, true, true, rwnd, null);
        sendPacket(finPacket);
        sendSeq++;
        
        // 等待确认
        while (true) {
            try {
                TCPacket tcpPacket = receiveQueue.poll(1000, TimeUnit.MILLISECONDS);
                if (tcpPacket == null) continue;
                
                if (tcpPacket.ackFlag && tcpPacket.ack == sendSeq && tcpPacket.fin) {
                    // 发送最后的ACK
                    TCPacket ackPacket = new TCPacket(localPort, remotePort, sendSeq, tcpPacket.seq + 1, 
                                                    false, true, false, rwnd, null);
                    sendPacket(ackPacket);
                    break;
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        
        isConnected = false;
        retransmitTimer.cancel();
    }
    
    // 处理接收到的数据包
    private void handlePacket(TCPacket packet) throws IOException {
    	System.out.println("接收数据包：" + packet);
    	System.out.println("received " + packet.data.length + " bytes: " + new String(packet.data));
        if (packet.fin) {
            // 处理结束请求
            recvSeq = packet.seq + 1;
            TCPacket ackPacket = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                            false, true, true, rwnd, null);
            sendPacket(ackPacket);
            isClosing = true;
            isConnected = false;
            return;
        }
        
        // 更新接收窗口
        rwnd = packet.window;
        
        // 处理确认
        if (packet.ackFlag && packet.ack > lastAck) {
            // 移除已确认的包
            synchronized (unackedPackets) {
                List<Long> toRemove = new ArrayList<>();
                for (long seq : unackedPackets.keySet()) {
                    TCPacket p = unackedPackets.get(seq);
                    int packetLength = (p.data != null ? p.data.length : 0) + 
                                      (p.syn || p.fin ? 1 : 0); // SYN和FIN包占用一个序列号
                    
                    if (seq + packetLength <= packet.ack) {
                        toRemove.add(seq);
                    }
                }
                toRemove.forEach(unackedPackets::remove);
            }
            
            lastAck = packet.ack;
            
            // 拥塞控制
            if (cwnd < ssthresh) {
                cwnd++;  // 慢启动
            } else {
                cwnd += 1.0 / cwnd;  // 拥塞避免
            }
        }
        
        // 处理数据
        if (packet.data != null && packet.data.length > 0) {
            if (packet.seq == recvSeq) {
                // 期望的数据，直接放入接收队列
                receiveBuffer.put(packet.seq, packet.data);
                processReceiveBuffer();
                // 跟processReceiveBuffer()加重复了
                //recvSeq += packet.data.length;
            } else if (packet.seq > recvSeq) {
                // 乱序到达的数据，先缓存
                receiveBuffer.put(packet.seq, packet.data);
            }
        }
    }
    
    // 处理接收缓冲区，将连续的数据拼接
    private void processReceiveBuffer() {
        List<byte[]> dataChunks = new ArrayList<>();
        
        while (receiveBuffer.containsKey(recvSeq)) {
            dataChunks.add(receiveBuffer.remove(recvSeq));
            recvSeq += dataChunks.get(dataChunks.size() - 1).length;
        }
        
        if (!dataChunks.isEmpty()) {
            // 拼接所有连续数据
            int totalLength = dataChunks.stream().mapToInt(b -> b.length).sum();
            byte[] allData = new byte[totalLength];
            int offset = 0;
            for (byte[] chunk : dataChunks) {
                System.arraycopy(chunk, 0, allData, offset, chunk.length);
                offset += chunk.length;
            }
            
            try {
            	// 发送确认
                sendAck();
            	
                appReceiveQueue.put(allData);
            } catch (Exception e) {
                Thread.currentThread().interrupt();
            }
        }
    }
    
    // 发送确认包
    private void sendAck() throws IOException {
        TCPacket ackPacket = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                        false, true, false, rwnd, null);
        System.out.println("回复ack: " + ackPacket);
        sendPacket(ackPacket);
    }
    
    // 发送数据包
    private void sendPacket(TCPacket packet) throws IOException {
        byte[] data = packet.toBytes();
        DatagramPacket datagram = new DatagramPacket(data, data.length, remoteAddr, remotePort);
        socket.send(datagram);
        
        // 需要重传的包：数据报、SYN包、FIN包
        if ((packet.data != null && packet.data.length > 0) || packet.syn || packet.fin) {
            synchronized (unackedPackets) {
                if (!unackedPackets.containsKey(packet.seq)) {
                    unackedPackets.put(packet.seq, packet);
                    
                    // 设置重传定时器
                    retransmitTimer.schedule(new TimerTask() {
                        @Override
                        public void run() {
                            synchronized (unackedPackets) {
                                if (unackedPackets.containsKey(packet.seq) && !isClosing) {
                                    try {
                                        // 重传
                                        TCPacket retransmitPacket = unackedPackets.get(packet.seq);
                                        byte[] retransmitData = retransmitPacket.toBytes();
                                        DatagramPacket retransmitDatagram = new DatagramPacket(
                                            retransmitData, retransmitData.length, remoteAddr, remotePort);
                                        socket.send(retransmitDatagram);
                                        
                                        // 拥塞控制：超时重传调整
                                        ssthresh = Math.max(cwnd / 2, 1);
                                        cwnd = 1;
                                    } catch (IOException e) {
                                        e.printStackTrace();
                                    }
                                }
                            }
                        }
                    }, 10000000); // 1秒后重传
                }
            }
        }
    }
    
    // 启动发送线程
    private void startSender() {
        new Thread(() -> {
            try {
                while (isConnected && !isClosing) {
                    // 根据拥塞窗口和接收窗口决定可以发送的数据量
                    int maxSegments = (int) Math.min(cwnd, rwnd / mss);
                    int segmentsSent = 0;
                    
                    while (segmentsSent < maxSegments && !sendQueue.isEmpty()) {
                        byte[] data = sendQueue.poll(100, TimeUnit.MILLISECONDS);
                        if (data == null) break;
                        
                        TCPacket packet = new TCPacket(localPort, remotePort, sendSeq, recvSeq, 
                                                    false, true, false, rwnd, data);
                        System.out.println("发送业务数据：" + packet);
                        sendPacket(packet);
                        
                        sendSeq += data.length;
                        segmentsSent++;
                    }
                    
                    Thread.sleep(10);
                }
            } catch (Exception e) {
                if (!isClosing) {
                    e.printStackTrace();
                }
            }
        }).start();
    }
    
    // 启动接收处理线程（处理分发过来的数据）
    private void startReceiverProcessor() {
        new Thread(() -> {
            try {
                while (isConnected && !isClosing) {
                    // 从receiveQueue获取分发的数据包
                    TCPacket tcpPacket = receiveQueue.poll(100, TimeUnit.MILLISECONDS);
                    if (tcpPacket != null) {
                        handlePacket(tcpPacket);
                    }
                }
            } catch (Exception e) {
                if (!isClosing) {
                    e.printStackTrace();
                }
            }
        }).start();
    }
    
    // 由服务器调用，分发数据包到当前连接
    public void dispatchPacket(TCPacket packet) {
        try {
            receiveQueue.put(packet);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
    
    public boolean isConnected() {
        return isConnected;
    }
}